import pymysql
import uuid
from loguru import logger
from config import MYSQL_HOST, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE, MYSQL_PORT


# 创建一个同步MySQL处理器
class MySQLHandler:
    def __init__(self, host, user, password, database, port=3306):
        self.db_config = {
            'host': host,
            'user': user,
            'password': password,
            'db': database,
            'port': port,
            'charset': 'utf8mb4',
            'autocommit': True
        }

        # 调用创建数据库的方法
        self.create_database_if_not_exists()

    def get_connection(self):
        """
        获取数据库连接
        :return: 返回一个连接对象
        """
        try:
            conn = pymysql.connect(**self.db_config, cursorclass=pymysql.cursors.DictCursor)
            return conn
        except Exception as e:
            logger.error(f"无法连接到数据库: {str(e)}")
            raise

    def create_database_if_not_exists(self):
        """
        检查数据库是否存在，如果不存在则创建
        """
        conn = pymysql.connect(host=self.db_config['host'], user=self.db_config['user'],
                               password=self.db_config['password'])
        with conn.cursor() as cursor:
            try:
                cursor.execute(f"CREATE DATABASE IF NOT EXISTS {self.db_config['db']}")
                logger.debug(f"数据库 {self.db_config['db']} 已创建或已存在")
            except Exception as e:
                logger.error(f"创建数据库失败: {str(e)}")
                raise
            finally:
                conn.close()

    def create_table_if_not_exists(self):
        """
        检查表是否存在，如果不存在则创建表
        """
        create_table_query = """
        CREATE TABLE IF NOT EXISTS voice_clone_tasks (
            id INT AUTO_INCREMENT PRIMARY KEY,
            uid VARCHAR(255) NOT NULL,
            timbre_id VARCHAR(255) NOT NULL UNIQUE,
            url TEXT NOT NULL,
            task_id VARCHAR(255) NOT NULL UNIQUE,
            task_status VARCHAR(50) DEFAULT 'pending',  -- 任务状态，默认为 'pending'
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            try:
                cursor.execute(create_table_query)
                logger.debug(f"表 voice_clone_tasks 已创建或已存在")
            except Exception as e:
                logger.error(f"创建表失败: {str(e)}")
                raise
            finally:
                conn.close()

    def execute_query(self, query, params=None):
        """
        执行查询语句（SELECT）
        :param query: SQL查询语句
        :param params: 查询参数
        :return: 查询结果
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            try:
                cursor.execute(query, params)
                result = cursor.fetchall()
                return result
            except Exception as e:
                logger.error(f"查询执行失败: {str(e)}")
                raise
            finally:
                conn.close()

    def execute_non_query(self, query, params=None):
        """
        执行非查询语句（INSERT, UPDATE, DELETE）
        :param query: SQL语句
        :param params: 参数
        :return: 影响的行数
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            try:
                cursor.execute(query, params)
                affected_rows = cursor.rowcount
                return affected_rows
            except Exception as e:
                logger.error(f"非查询执行失败: {str(e)}")
                raise
            finally:
                conn.close()

    def insert_voice_clone_task(self, uid: str, timbre_id: str, url: str, task_id: str,
                                task_status: str = "pending"):
        """
        Args:
            uid:
            timbre_id:
            url:
            task_id:
            task_status:

        Returns:

        """
        # 确保表已存在
        self.create_table_if_not_exists()

        query = """
        INSERT INTO voice_clone_tasks (uid, timbre_id, url, task_id, task_status)
        VALUES (%s, %s, %s, %s, %s)
        """
        params = (uid, timbre_id, url, task_id, task_status)
        try:
            self.execute_non_query(query, params)
            logger.debug(f"任务 {task_id} 成功插入数据库，状态为 {task_status}")
        except Exception as e:
            logger.error(f"插入任务信息到数据库时失败: {str(e)}")
            raise

    def get_voice_clone_task(self, timbre_id: str):
        """
        Args:
            timbre_id:

        Returns:

        """
        query = "SELECT URL FROM voice_clone_tasks WHERE timbre_id = %s"
        try:
            result = self.execute_query(query, (timbre_id,))
            return result
        except Exception as e:
            logger.error(f"查询音色 {timbre_id} 失败: {str(e)}")
            raise

    # 通过任务ID获取任务
    def get_voice_clone_task_by_task_id(self, task_id: str):
        """
        Args:
            task_id:

        Returns:

        """
        query = "SELECT * FROM voice_clone_tasks WHERE task_id = %s"
        try:
            result = self.execute_query(query, (task_id,))
            return result
        except Exception as e:
            logger.error(f"查询任务 {task_id} 失败: {str(e)}")
            raise

    def update_voice_clone_task_status(self, task_id: str, new_status: str):
        """
        Args:
            task_id:
            new_status:

        Returns:

        """
        query = "UPDATE voice_clone_tasks SET task_status = %s WHERE task_id = %s"
        params = (new_status, task_id)
        try:
            affected_rows = self.execute_non_query(query, params)
            if affected_rows == 0:
                logger.warning(f"任务 {task_id} 未找到，更新失败")
            else:
                logger.debug(f"任务 {task_id} 的状态成功更新为 {new_status}")
        except Exception as e:
            logger.error(f"更新任务 {task_id} 状态失败: {str(e)}")
            raise


# 创建一个同步MySQL处理器
mysql_handler = MySQLHandler(host=MYSQL_HOST, user=MYSQL_USER, password=MYSQL_PASSWORD,
                             database=MYSQL_DATABASE, port=MYSQL_PORT)
